﻿///////////////////////////////////////////////////////////////////////////////
/* Code Usage: Parts of the code in this class has been referenced from
 * the MSDN article by Scott Mitchell at
 * http://msdn.microsoft.com/en-us/library/ms379573(VS.80).aspx
 *  Especially the traversals.
*/
///////////////////////////////////////////////////////////////////////////////

using System;
using System.Collections;
using System.Collections.Generic;

namespace ComputationalGeometry.DataStructures
{
    /// <summary>
    /// Binary Search Tree.
    /// </summary>
    /// <typeparam name="T">The data the tree will hold.</typeparam>
    public class BinaryTree<T> : ICollection<T>, IEnumerable<T>
    {
        private BinaryTreeNode<T> _root = null;
        private int _count = 0;
        private IComparer<T> _comparer = Comparer<T>.Default;

        public BinaryTree()
        {
        }

        public BinaryTree(IComparer<T> comparer)
        {
            _comparer = comparer;
        }

        public IComparer<T> DataComparer
        {
            get { return _comparer; }
            set { _comparer = value; }
        }

        public int Count
        {
            get { return this._count; }
        }

        public bool IsReadOnly
        {
            get { return false; }
        }

        public void Add(T item)
        {
            BinaryTreeNode<T> newNode = new BinaryTreeNode<T>(item);
            BinaryTreeNode<T> cur = _root;
            BinaryTreeNode<T> prev = null;
            int result = 0;

            while (cur != null)
            {
                prev = cur;

                result = _comparer.Compare(cur.Value, item);
                cur = result < 0 ? cur.Right : cur.Left;
            }

            if (_root == null)
            {
                _root = new BinaryTreeNode<T>(item);
            }
            else
            {
                if (prev != null)
                {
                    result = _comparer.Compare(prev.Value, item);
                    if (result < 0)
                    {
                        // add to right child
                        prev.Right = newNode;
                    }
                    else
                    {
                        // else, add to left child
                        prev.Left = newNode;
                    }
                }
            }

            _count++;
        }

        public void Clear()
        {
            _root = null;
            _count = 0;
        }

        public bool Contains(T item)
        {
            if (_root != null)
            {
                BinaryTreeNode<T> cur = _root;
                while (cur != null)
                {
                    int result = _comparer.Compare(cur.Value, item);
                    if (result == 0)
                    {
                        return true;
                    }
                    else if (result > 0)
                    {
                        cur = cur.Left;
                    }
                    else if (result < 0)
                    {
                        cur = cur.Right;
                    }
                }
            }

            return false;
        }

        public void CopyTo(T[] array, int arrayIndex)
        {
            CopyTo(array, arrayIndex, TreeTraversalMethod.Inorder);
        }

        public void CopyTo(T[] array, int index, TreeTraversalMethod traversalMethod)
        {
            IEnumerable<T> enumProp = null;

            switch (traversalMethod)
            {
                case TreeTraversalMethod.Preorder:
                    enumProp = Preorder;
                    break;

                case TreeTraversalMethod.Inorder:
                    enumProp = Inorder;
                    break;

                default:
                    enumProp = Postorder;
                    break;
            }

            int i = 0;
            foreach (T value in enumProp)
            {
                array[i + index] = value;
                i++;
            }
        }

        public bool Remove(T item)
        {
            // check if tree is empty
            if (_root == null)
            {
                return false;
            }

            BinaryTreeNode<T> cur = this._root;
            BinaryTreeNode<T> prev = null;
            bool nodeFound = false;

            // find the node to remove
            while (cur != null && nodeFound != true)
            {
                int compareResult = _comparer.Compare(cur.Value, item);
                if(compareResult == 0)
                {
                    nodeFound = true;
                }
                else
                {
                    prev = cur;

                    cur = compareResult > 0 ? cur.Left : cur.Right;
                }
            }

            // if we didn't find the node, return false
            if (!nodeFound)
            {
                return false;
            }

            // if we did find the node, there are 3 removal cases.
            if (cur.Left == null)
            {
                // Case 1: node has empty left subtree and non-empty right subtree
                // (or leaf node)
                if (prev == null)
                {
                    _root = cur.Right;
                }
                else if (prev.Left == cur)
                {
                    prev.Left = cur.Right;
                }
                else
                {
                    prev.Right = cur.Right;
                }
            }
            else if (cur.Right == null)
            {
                // Case 2: node has empty right subtree and a non-empty left
                if (prev == null)
                {
                    _root = cur.Left;
                }
                if (prev.Left == cur)
                {
                    prev.Left = cur.Left;
                }
                else
                {
                    prev.Right = cur.Left;
                }
            }
            else
            {
                // case 3: node has both children
                BinaryTreeNode<T> start = cur;
                cur = start.Left;
                prev = null;

                // find right-most left child
                while (cur.Right != null)
                {
                    prev = cur;
                    cur = cur.Right;
                }

                // exchange data
                start.Value = cur.Value;

                // now apply case 2
                if (prev == null)
                {
                    start.Left = cur.Left;
                }
                else
                {
                    prev.Right = cur.Left;
                }
            }

            // remove the node
            cur.Left = null;
            cur.Right = null;
            cur = null;

            // update node count
            _count--;

            return true;
        }

        IEnumerator<T> IEnumerable<T>.GetEnumerator()
        {
            return GetEnumerator(TreeTraversalMethod.Inorder);
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator(TreeTraversalMethod.Inorder);
        }
        
        public virtual IEnumerator<T> GetEnumerator(TreeTraversalMethod traversalMethod)
        {
            switch (traversalMethod)
            {
                case TreeTraversalMethod.Preorder:
                    return Preorder.GetEnumerator();

                case TreeTraversalMethod.Inorder:
                    return Inorder.GetEnumerator();

                default:
                    return Postorder.GetEnumerator();
            }
        }

        public IEnumerable<T> Preorder
        {
            get
            {
                // A single stack is sufficient here - it simply maintains the correct
                // order with which to process the children.
                Stack<BinaryTreeNode<T>> toVisit = new Stack<BinaryTreeNode<T>>(Count);
                BinaryTreeNode<T> current = _root;
                if (current != null) toVisit.Push(current);

                while (toVisit.Count != 0)
                {
                    // take the top item from the stack
                    current = toVisit.Pop();

                    // add the right and left children, if not null
                    if (current.Right != null) toVisit.Push(current.Right);
                    if (current.Left != null) toVisit.Push(current.Left);

                    // return the current node
                    yield return current.Value;
                }
            }
        }

        public IEnumerable<T> Inorder
        {
            get
            {
                // A single stack is sufficient - this code was made available by Grant Richins:
                // http://blogs.msdn.com/grantri/archive/2004/04/08/110165.aspx
                Stack<BinaryTreeNode<T>> toVisit = new Stack<BinaryTreeNode<T>>(Count);
                for (BinaryTreeNode<T> current = _root; current != null || toVisit.Count != 0; current = current.Right)
                {
                    // Get the left-most item in the subtree, remembering the path taken
                    while (current != null)
                    {
                        toVisit.Push(current);
                        current = current.Left;
                    }

                    current = toVisit.Pop();
                    yield return current.Value;
                }
            }
        }

        public IEnumerable<T> Postorder
        {
            get
            {
                // maintain two stacks, one of a list of nodes to visit,
                // and one of booleans, indicating if the note has been processed
                // or not.
                Stack<BinaryTreeNode<T>> toVisit = new Stack<BinaryTreeNode<T>>(Count);
                Stack<bool> hasBeenProcessed = new Stack<bool>(Count);
                BinaryTreeNode<T> current = _root;
                if (current != null)
                {
                    toVisit.Push(current);
                    hasBeenProcessed.Push(false);
                    current = current.Left;
                }

                while (toVisit.Count != 0)
                {
                    if (current != null)
                    {
                        // add this node to the stack with a false processed value
                        toVisit.Push(current);
                        hasBeenProcessed.Push(false);
                        current = current.Left;
                    }
                    else
                    {
                        // see if the node on the stack has been processed
                        bool processed = hasBeenProcessed.Pop();
                        BinaryTreeNode<T> node = toVisit.Pop();
                        if (!processed)
                        {
                            // if it's not been processed, "recurse" down the right subtree
                            toVisit.Push(node);
                            hasBeenProcessed.Push(true);    // it's now been processed
                            current = node.Right;
                        }
                        else
                            yield return node.Value;
                    }
                }
            }
        }

        public int Height()
        {
            return Height(_root);
        }

        private static int Height(BinaryTreeNode<T> root)
        {
            if (root == null)
            {
                return 0;
            }

            int rightHeight = Height(root.Right);
            int leftHeight = Height(root.Left);

            return Math.Max(rightHeight, leftHeight) + 1;
        }
    }    
}
